from repository.SampleRepository import SampleRepository
from repository.FeatureRepository import FeatureRepository
from HGMSim.HGMSimModelWrapper import HGMSimModelWrapper
from basicTools.myTLSH import MyTlsh


class SimilarityService:
    def __init__(self):
        self.sampleRepository = SampleRepository()
        self.featureRepository = FeatureRepository()
        self.myTlsh = MyTlsh()
        # 加载之前训练好的模型
        self.hGMSimModelWrapper = HGMSimModelWrapper()

    def pairSim(self, sampleId1Md5, sampleId2Md5):
        """
        计算两个样本对之间的相似值
        """
        return self.hGMSimModelWrapper.calculateSim([(sampleId1Md5, sampleId2Md5, 1)])

    def simTopK(self, sampleIdMd5, kValue, returnKValue):
        """
        sampleIdMd: target sample's md5 value
        kValue: select top k samples
        """
        # step1: 先粗筛出2k个样本
        names = self.findSampleNameByEmbedding(sampleIdMd5, kValue)
        # step2: 根据初筛的样本筛选出k个样本
        dataInf = []
        for name in names:
            dataInf.append((sampleIdMd5, name, 1))
        simScores = self.hGMSimModelWrapper.calculateSim(dataInf)
        ans = []
        for i in range(len(dataInf)):
            ans.append((dataInf[i][1], simScores[i]))

        # step3: 构建结果数据
        res = sorted(ans, key=lambda x: -x[1])[:returnKValue]
        print("相似性度量检索的结果为：")
        print(res)
        return res

    def findSampleNameByEmbedding(self, sampleIdMd5, kValue):
        """
        calculate the top 100 samples for similarity by tlsh value
        """
        # obtain target sample's tlsh
        tarEmbedding = self.featureRepository.queryEmbeddingByMd5(sampleIdMd5)
        # obtain all samples' tlsh values
        allSampleEmbeddings = self.featureRepository.queryAllEmbedding()
        # calculate the topk samples
        ans = self.myTlsh.getTopkscore(tarEmbedding, allSampleEmbeddings, kValue)
        return ans

    def querySimList(self, operatorName):
        """
        请求某样本的相似性度量
        """
        pass

    def queryPairSimInf(self, sampleId1, sampleId2):
        """
        获得两个样本的相似性度量对比结果
        """
        pass
